﻿#' @TODO 药物IC50相关分析
#' @title 药物IC50相关分析
#' @description 利用计算得到的各个药物在样本中的IC50数据，计算相关性，以及组间是否显著差异,生成相关图片.
#' @param DrugInfor 药物IC50信息，包含样本列与药物信息
#' @param GroupInfor 用于组间样本秩和检验，包含样本列`sample`，比较同药物在不同样本分组的IC50是否有显著性差异，如果为`NULL`则不进行该分析步骤
#' @param Df_For_Corr 数据框，包含样本列`sample`，与其他信息，用于计算与药物之间的spearman相关性，
#' 可以是模型基因表达量或者riskscore等连续变量，如果为`NULL`则不进行该分析步骤
#' @param saveFile 是否保存文件，默认FALSE
#' @param output_dir 文件输出路径
#' @param DrugSelected 需要展示的药物，如果为NULL则不进行绘图，根据输入的数据会绘制分组间boxplot；连续变量的相关性散点图。
#' @param y_lab_chr 字符串，图中纵坐标轴字段，如果为NULL，默认为log(IC50)
#' @param w_single 单个图片宽度
#' @param h_single 单个图片高度
#' @usage 数据以及运用示例
#'
#' #数据准备
#' library(tidyverse)
#' load("/Pub/Users/wangyk/project/tmp/mimandaB/A_vs_BC/major.RData",verbose = T)
#' > dim(training[[3]])
#' [1] 470   5
#'
#' drug_data <- data.table::fread("/Pub/Users/wangyk/project/tmp/mimandaB/train_expr.drug_sensitivity_predicted.txt") %>%
#'     as.data.frame() %>%
#'     dplyr::rename(sample = 1)
#' > drug_data[1:4,1:4]
#'      sample   A.443654 A.770041  ABT.263
#' 1 GSM775979 -0.5689866 4.022199 2.504680
#' 2 GSM775980 -0.5745180 4.004569 2.387414
#' 3 GSM775981 -0.5686237 3.860390 2.529512
#' 4 GSM775982 -0.5575950 3.978383 2.799947
#'
#' aa <- training[[3]] %>% select(sample,riskgroup)
#' cc <- training[[3]] %>% select(sample,riskscore,time)
#' aa$cluster <- c(rep('A',100),rep('B',370))
#' aa$cluster2 <- c(rep('A',200),rep('B',270))
#' aa$riskgroup[1:6] <- NA
#' > head(aa)
#'      sample riskgroup cluster cluster2
#' 1 GSM775979      <NA>       A        A
#' 2 GSM775980      <NA>       A        A
#' 3 GSM775981      <NA>       A        A
#' 4 GSM775982      <NA>       A        A
#' 5 GSM775983      <NA>       A        A
#' 6 GSM775984      <NA>       A        A
#' > head(cc)
#'      sample  riskscore    time
#' 1 GSM775979 -0.5474144  61.897
#' 2 GSM775980 -0.5143293  90.743
#' 3 GSM775981 -1.2046581 102.275
#' 4 GSM775982 -0.9046476 102.834
#' 5 GSM775983 -0.3226878  82.990
#' 6 GSM775984 -1.2695206  50.563
#'
#' #运用示例
#' b <- drug_analyse(DrugInfor = drug_data, GroupInfor = aa, Df_For_Corr = cc, saveFile = F, output_dir = "./")
#' > names(b)
#' [1] "Wli_res"  "corr_res"
#'
#' @return *list*
#' @export
#' @author *WYK*
#'
Drug_IC50 <- function(DrugInfor = NULL, GroupInfor = NULL, Df_For_Corr = NULL,
                      saveFile = T, output_dir = "./", DrugSelected = NULL, y_lab_chr = NULL,
                      w_single = 2.5, h_single = 2.7) {
    library(tidyverse)
    library(psych)

    if (is.null(y_lab_chr)) {
        y_lab_chr <- "log(IC50)"
    }

    if (!is.null(GroupInfor)) {
        if (dim(GroupInfor)[2] > 2) {
            Drug_name <- DrugInfor %>%
                dplyr::select(-sample) %>%
                colnames()
            group_names <- base::setdiff(colnames(GroupInfor), "sample")

            Wli_res <- map_dfc(group_names, function(group_name) {
                p_val_drug_in_group <- map_dbl(Drug_name, function(x) {
                    df_tmp <- inner_join(DrugInfor %>% dplyr::select(sample, x), GroupInfor)
                    wilcox_res <- wilcox.test(df_tmp %>% pull(x) ~ get0(group_name), data = df_tmp)
                    return(wilcox_res$p.value)
                })

                p.signif <- cut(p_val_drug_in_group, breaks = c(0, 0.0001, 0.001, 0.01, 0.05, 1), labels = c("****", "***", "**", "*", " "))

                Wli_res <- tibble(
                    "Drug_name" = Drug_name,
                    "group" = p_val_drug_in_group,
                    "p.signif" = p.signif
                )

                colnames(Wli_res)[2] <- paste0("WilcoxTsetPval_in_", group_name)
                return(Wli_res)
            })
        } else {
            Drug_name <- DrugInfor %>%
                select(-sample) %>%
                colnames()
            group_name <- base::setdiff(colnames(GroupInfor), "sample")

            p_val_drug_in_group <- map_dbl(Drug_name, function(x) {
                df_tmp <- inner_join(DrugInfor %>% dplyr::select(sample, x), GroupInfor)
                wilcox_res <- wilcox.test(df_tmp %>% pull(x) ~ get0(group_name), data = df_tmp)
                return(wilcox_res$p.value)
            })

            p.signif <- cut(p_val_drug_in_group, breaks = c(0, 0.0001, 0.001, 0.01, 0.05, 1), labels = c("****", "***", "**", "*", " "))

            Wli_res <- tibble(
                "Drug_name" = Drug_name,
                "group" = p_val_drug_in_group,
                "p.signif" = p.signif
            )

            colnames(Wli_res)[2] <- paste0("WilcoxTsetPval_in_", group_name)
        }
    }

    if (!is.null(Df_For_Corr)) {
        data_for_corr <- inner_join(DrugInfor, Df_For_Corr)
        varible_name <- Df_For_Corr %>%
            select(-sample) %>%
            colnames()
        drug_name <- DrugInfor %>%
            select(-sample) %>%
            colnames()

        x_df <- data_for_corr %>%
            dplyr::select(all_of(drug_name)) %>%
            as.matrix()
        y_df <- data_for_corr %>%
            dplyr::select(all_of(varible_name)) %>%
            as.matrix()

        corr_res <- corr.test(x = x_df, y = y_df, method = "spearman", adjust = "BH")
    }

    if (!is.null(DrugSelected) && !is.null(GroupInfor)) {

        # 判断总共有图片排队几行
        N <- length(DrugSelected)
        nrow_used <- case_when(
            between(N, 1, 4) ~ 1,
            between(N, 5, 8) ~ 2,
            between(N, 9, 12) ~ 3,
            between(N, 13, 16) ~ 4,
            between(N, 17, 20) ~ 5,
            N >= 21 ~ 6
        )

        width_used <- case_when(
            N < 5 ~ case_when(
                N == 1 ~ w_single * 1,
                N == 2 ~ w_single * 2,
                N == 3 ~ w_single * 3,
                N == 4 ~ w_single * 4
            ),
            N > 4 ~ w_single * 4
        )
        # 根据图片排列行数 判断出图高度
        height_used <- switch(nrow_used,
            h_single,
            h_single * 2,
            h_single * 3,
            h_single * 4,
            h_single * 5,
            h_single * 6
        )

        set.seed(1110)

        group_plot_list <- map(colnames(GroupInfor %>% select(-sample)), function(x) {
            #  x <- 'riskgroup'
            GroupInfor_1 <- GroupInfor %>% dplyr::select("sample", x)

            plot_df <- DrugInfor %>%
                dplyr::select("sample", DrugSelected) %>%
                inner_join(GroupInfor_1) %>%
                pivot_longer(col = -c(sample, x), names_to = "DrugSelected", values_to = "IC50")

            plot_df$DrugSelected <- factor(plot_df$DrugSelected, levels = DrugSelected)

            group_plot <- plot_df %>%
                ggplot(aes(x = factor(get0(x)), y = IC50)) +
                geom_boxplot(aes(fill = get0(x)),
                    color = "black",
                    width = .5, show.legend = F, outlier.size = .55, alpha = .9, size = .6
                ) +
                # geom_boxplot(aes(), width = .4, show.legend = F, outlier.shape = NA, alpha = 1) +
                # geom_jitter(aes(color = get0(x)), width = .2, show.legend = F, alpha = .65) +
                scale_fill_manual(values = RColorBrewer::brewer.pal(9, "Set1")[c(2, 1, 3:9)]) +
                labs(x = x, y = y_lab_chr) +
                theme(legend.position = NULL) +
                ggpubr::theme_pubr() +
                ggpubr::stat_compare_means(label.x.npc = "left", label.y.npc = "top") +
                facet_wrap(. ~ DrugSelected, scales = "free", nrow = nrow_used)+
                coord_cartesian(clip = "off")

            return(group_plot)
        })
        names(group_plot_list) <- colnames(GroupInfor %>% dplyr::select(-sample))
    }

    if (!is.null(DrugSelected) && !is.null(Df_For_Corr)) {
        N <- length(DrugSelected)
        nrow_used <- case_when(
            between(N, 1, 4) ~ 1,
            between(N, 5, 8) ~ 2,
            between(N, 9, 12) ~ 3,
            between(N, 13, 16) ~ 4,
            between(N, 17, 20) ~ 5,
            N >= 21 ~ 6
        )

        width_used_cor <- case_when(
            N < 5 ~ case_when(
                N == 1 ~ w_single * 1,
                N == 2 ~ w_single * 2,
                N == 3 ~ w_single * 3,
                N == 4 ~ w_single * 4
            ),
            N > 4 ~ w_single * 4
        )
        # 根据图片排列行数 判断出图高度
        height_used_cor <- switch(nrow_used,
            h_single,
            h_single * 2,
            h_single * 3,
            h_single * 4,
            h_single * 5,
            h_single * 6
        )

        corPoint_plot_list <- map(colnames(Df_For_Corr %>% dplyr::select(-sample)), function(x) {
            Df_For_Corr_1 <- Df_For_Corr %>%
                dplyr::select("sample", x) %>%
                arrange(get(x))

            single_plot <- map(DrugSelected, function(y) {
                corPoint_plot <- DrugInfor %>%
                    dplyr::select("sample", y) %>%
                    inner_join(Df_For_Corr_1) %>%
                    pivot_longer(col = -c(sample, x), names_to = "DrugSelected", values_to = "IC50") %>%
                    ggplot(aes(x = get0(x), y = IC50)) +
                    geom_point(color = "#377EB8", alpha = .9) +
                    geom_smooth(method = "lm", col = "#E41A1C", alpha = .55) +
                    ggpubr::stat_cor(method = "spearman") +
                    labs(x = x, y = y_lab_chr, title = y) +
                    egg::theme_article(15) +
                    theme(plot.title = element_text(hjust = 0.5))

                return(corPoint_plot)
            })

            names(single_plot) <- str_c(DrugSelected, "_in_", x)

            corPoint_plot <- cowplot::plot_grid(plotlist = single_plot, nrow = nrow_used)

            return(corPoint_plot)
        })

        names(corPoint_plot_list) <- colnames(Df_For_Corr %>% dplyr::select(-sample))

        corPoint_SinglePlot_list <- map(colnames(Df_For_Corr %>% dplyr::select(-sample)), function(x) {
            Df_For_Corr_1 <- Df_For_Corr %>%
                dplyr::select("sample", x) %>%
                arrange(get(x))

            single_plot <- map(DrugSelected, function(y) {
                corPoint_plot <- DrugInfor %>%
                    dplyr::select("sample", y) %>%
                    inner_join(Df_For_Corr_1) %>%
                    pivot_longer(col = -c(sample, x), names_to = "DrugSelected", values_to = "IC50") %>%
                    ggplot(aes(x = get0(x), y = IC50)) +
                    geom_point(color = "#377EB8", alpha = .9) +
                    geom_smooth(method = "lm", col = "#E41A1C", alpha = .55) +
                    ggpubr::stat_cor(method = "spearman") +
                    labs(x = x, y = y_lab_chr, title = y) +
                    egg::theme_article(15) +
                    theme(plot.title = element_text(hjust = 0.5))

                return(corPoint_plot)
            })

            names(single_plot) <- str_c(DrugSelected, "_in_", x)

            return(single_plot)
        })

        names(corPoint_SinglePlot_list) <- colnames(Df_For_Corr %>% dplyr::select(-sample))


        if (file.exists("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/一些数据/drug_db/GDSC1/drug_pathway.tsv")) {
            # drug_pathway <- read.delim("/Pub/Users/liulk/RCodes/RCodes_LLK/drug_susceptibility/DRUG_PATHWAY.txt")
            drug_pathway <- read.delim("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/一些数据/drug_db/GDSC1/drug_pathway.tsv") %>%
                select(1, 3) %>%
                rename(DRUG_NAME = 1, TARGET_PATHWAY = 2)

            drug_name <- DrugSelected

            corr_res_p <- corr_res$p %>%
                .[DrugSelected, , drop = F] %>%
                t()
            corr_res_r <- corr_res$r %>%
                .[DrugSelected, , drop = F] %>%
                t()

            corr_res_p_sig <- map_dfr(1:nrow(corr_res_p), function(x) {
                a <- map_dfc(1:ncol(corr_res_p), function(y) {
                    case_when(
                        between(corr_res_p[x, y], 0.01, 0.05) ~ "*",
                        between(corr_res_p[x, y], 0.001, 0.01) ~ "**",
                        between(corr_res_p[x, y], 0.0001, 0.001) ~ "***",
                        corr_res_p[x, y] < 0.0001 ~ "****",
                        corr_res_p[x, y] > 0.05 ~ " "
                    )
                }) %>% as.data.frame()
            }) %>% as.data.frame()

            rownames(corr_res_p_sig) <- rownames(corr_res_p)
            colnames(corr_res_p_sig) <- colnames(corr_res_p)

            library(ComplexHeatmap)
            library(RColorBrewer)
            # c("#377EB8", "white", "firebrick3")
            col_zscore <- circlize::colorRamp2(c(-1, 0, 1), c("#377EB8", "white", "firebrick3"))

            drug_pathway_infor <- data.frame(
                drug_name = DrugSelected,
                Target_Pathway = NA
            )

            drug_pathway_infor[, 2] <- map_chr(DrugSelected, function(x) {
                # x <- 'BI.2536_1086'
                x <- str_split(x, "_")[[1]][1] %>% str_replace_all(., pattern = "\\.", replacement = "_")

                if (x %in% drug_pathway$DRUG_NAME) {
                    pathway_infor <- drug_pathway %>%
                        filter(DRUG_NAME == x) %>%
                        pull(TARGET_PATHWAY) %>%
                        .[1]
                } else {
                    pathway_infor <- "Other"
                }
            })

            group_chara <- sort(drug_pathway_infor$Target_Pathway %>% unique())
            col_name <- c(brewer.pal(12, "Paired"), brewer.pal(12, "Paired"))
            col_name <- col_name[seq_along(group_chara)]
            names(col_name) <- group_chara

            if ("Other" %in% group_chara) {
                col_name[which(group_chara == "Other")] <- "#6b6b6b"
            }

            col_anno <- HeatmapAnnotation(
                df = drug_pathway_infor %>% rename(Target_Pathway = 2) %>% .[, 2, drop = F],
                col = list(
                    Target_Pathway = col_name
                ),
                annotation_legend_param = list(Target_Pathway = list(title = "Target\nPathway"))
            )

            lgd <- Legend(
                labels = c("ns", "<0.05", "<0.01", "<0.001", "<0.0001"),
                title = "P.val", size = .8,
                type = "points",
                background = "white",
                pch = c(" ", "*", "**", "***", "****")
            )

            height_heatmap <- if (nrow(corr_res_r) == 1) {
                times <- 1.5
            } else if (nrow(corr_res_r) %in% c(2:10)) {
                times <- 1
            } else {
                times <- .8
            }

            heatmap_res <- Heatmap(
                matrix = corr_res_r,
                name = "rho",
                col = col_zscore,
                show_column_dend = F,
                show_column_names = T,
                show_row_names = T,
                column_names_rot = 45,
                column_names_gp = gpar(fontsize = 11),
                cluster_rows = F,
                show_row_dend = F,
                rect_gp = gpar(col = "white", lwd = 1),
                # row_names_gp = gpar(fontsize = 8),
                column_title = " ",
                cell_fun = function(j, i, x, y, width, height, fill) {
                    if (corr_res_p_sig[i, j] != " ") {
                        grid.text(sprintf("%s", corr_res_p_sig[i, j]), x, y, gp = gpar(fontsize = 10), vjust = .05)
                        grid.text(sprintf("%.2f", corr_res_r[i, j]), x, y, gp = gpar(fontsize = 8), vjust = .9)
                    }
                },
                top_annotation = col_anno,
                column_split = drug_pathway_infor %>% dplyr::rename(Target_Pathway = 2) %>% .[, 2, drop = F],
                height = unit(times * nrow(corr_res_r), "cm"),
                border = F,
                width = unit(1.05 * length(DrugSelected), "cm"),
                # row_order = colnames(corr_res_r)
            )

            gg.heatmap_res <<- heatmap_res %>%
                draw(., merge_legend = F) %>% # annotation_legend_list = lgd
                grid.grabExpr() %>%
                ggplotify::as.ggplot() +
                theme(plot.margin = unit(c(.1, .5, .1, .7), "cm"))
        }
    }

    if (saveFile) {
        if (!dir.exists(str_glue("{output_dir}/Corr_infor/"))) {
            dir.create(str_glue("{output_dir}/Corr_infor/"), recursive = T)
        } else {
            message(str_glue("{output_dir}/Corr_infor/ is ready."))
        }

        if (!is.null(DrugSelected) && !is.null(GroupInfor)) {
            walk(seq_along(group_plot_list), function(x) {
                x <- 1
                plot_chara <- names(group_plot_list)[x]
                ggsave(
                    plot = group_plot_list[[x]], filename = str_glue("{output_dir}/Corr_infor/Figure_{plot_chara}_BoxPlot.pdf"),
                    width = width_used, height = height_used
                )
            })
        }

        if (!is.null(DrugSelected) && !is.null(Df_For_Corr)) {
            walk(seq_along(corPoint_plot_list), function(x) {
                # x <- 1
                plot_chara <- names(corPoint_plot_list)[x]
                ggsave(
                    plot = corPoint_plot_list[[x]], filename = str_glue("{output_dir}/Corr_infor/Figure_{plot_chara}_CorPlot.pdf"),
                    width = width_used_cor, height = height_used_cor
                )
            })

            walk(names(corPoint_SinglePlot_list), function(x) {
                # x <- names(corPoint_SinglePlot_list)[1]
                walk(names(corPoint_SinglePlot_list[[x]]), function(y) {
                    # y <- names(corPoint_SinglePlot_list[[x]])[1]
                    ggsave(
                        plot = corPoint_SinglePlot_list[[x]][[y]],
                        filename = str_glue("{output_dir}/Corr_infor/Figure_{y}_CorPlot.pdf"),
                        width = 3.3, height = 3.7
                    )
                })
            })

            # ggsave(
            #     plot = gg.heatmap_res, filename = str_glue("{output_dir}/Corr_infor/Figure_Drug_Heatmap.pdf"),
            #     width = 1.1*length(DrugSelected), height = (ncol(Df_For_Corr)-1)+1.5
            # )
            if (exists("gg.heatmap_res")) {
                ggsave(
                    plot = gg.heatmap_res, filename = str_glue("{output_dir}/Corr_infor/Figure_Drug_Heatmap.pdf"),
                    width = 1.2 * length(DrugSelected), height = (ncol(Df_For_Corr) - 1) + 1.5
                )
            }
        }

        if (exists("Wli_res")) {
            write_tsv(x = Wli_res, file = str_glue("{output_dir}/Corr_infor/Wli_res_pvalSig_in_Group.txt"))
        }

        if (exists("corr_res")) {
            write_tsv(
                x = corr_res$r %>% as.data.frame() %>% rownames_to_column("Drug"),
                file = str_glue("{output_dir}/Corr_infor/rho.txt")
            )
            write_tsv(
                x = corr_res$p %>% as.data.frame() %>% rownames_to_column("Drug"),
                file = str_glue("{output_dir}/Corr_infor/rho_pval.txt")
            )
        }
    }

    res <- map(c("Wli_res", "corr_res", "gg.heatmap_res"), function(x) {
        if (exists(x)) {
            return(get0(x))
        } else {
            NA
        }
    })

    names(res) <- c("Wli_res", "corr_res", "gg.heatmap_res")
    return(res)
}


logIC50Analyse <- function(DrugInfor = NULL, GroupInfor = NULL, Df_For_Corr = NULL,
                           saveFile = T, output_dir = "./", DrugSelected = NULL, y_lab_chr = NULL) {
    warning(stringr::str_glue('Function {crayon::blue("logIC50Analyse")} is replaced by {crayon::blue("Drug_IC50")},and it is out of maintenance.'))

    res <- Drug_IC50(
        DrugInfor = DrugInfor, GroupInfor = GroupInfor, Df_For_Corr = Df_For_Corr,
        saveFile = saveFile, output_dir = output_dir, DrugSelected = DrugSelected, y_lab_chr = y_lab_chr
    )

    return(res)
}



